% coreged.m - find "best" affine transform(s) to fit elp to scalp
%
% using Steven Michael's kdtree routine from:
% http://www.mathworks.com/matlabcentral/fileexchange/
%  7030-kd-tree-nearest-neighbor-and-range-search
%
% 1) load MRI and elp
% 2) extract scalp as vertex list
% 3) build k-d tree for nearest point searc
% 4) find initial closest points
% 5) display initial fit
% 6) use "fmincon", "lsqnonlin" or "fminsearch" to find best xform
% 7) display transformed fit
% 8) write transformed elp+acpc.BRIK & HEAD

function coreged(subj,date,file_type)

global MRI points_xyz TreeRoot ClosestPts TRxyz dd MR_th flags
global i_ter trackx0 track_dist output x1 fval options fit_time

% plt_flg=1; % 1 to plot, 0 not !!!use flags.plot !!!

in_path='c:/data/cap/';
s_fold=['S' subj '_' date '/'];
in_file=['S' subj file_type];

n_points=0;    % # of individual elp points (not continuous scan)

dd=16;  % decimation factor for vertex display (64 is fast but sparse
        % 16 is slow and dense)

scrsz = get(0,'ScreenSize');

% 1) load vertex (scalp) and (tai aligned) elp files
%
% Load Chris' scalp file
%
vert_file=[in_path s_fold 'S' subj '_' date '_IZI_scalp.mat'];
disp(['Loading vertex file ' vert_file])
load(vert_file,'-mat');
vert1=xyz_mask;

% load a tai aligned elp file
%
xyz_file=[in_path s_fold in_file '.mat'];
disp(['Loading ' xyz_file]);
load(xyz_file,'-mat');
if n_points==0;n_points=length(xyz);end
points_xyz=xyz(1:n_points,:);
% scatter3(points_xyz(:,1),points_xyz(:,2),points_xyz(:,3),'filled')
% axis equal
% axis off

% 3) Build the k-d Tree once from the reference datapoints.
%
disp('Building k-d tree')
TreeRoot=kdtree(vert1);

% 4) find initial closest points
%
disp('finding Closest Points');
[Pts_idx,ClosestPts] = kdtree_closestpoint(TreeRoot,points_xyz);

% 6) use minimization to find best xform
%    N.B. might want to use "lsqnonlin"
%    or other routines from "Minimization" toolbox
%
% disp('Minimizing total distance')

A= [];b=[];Aeq=[];beq=[];

if flags.params == 3  % 3 params (scale)
    x0=[1,1,1]; % start with null transform: shift;rot;scale
    x2=x0';
    x2=x2(:)';
    lb=[.8 .8 .8 ];
    ub=[1.2 1.2 1.2];
end

if flags.params == 6  % 6 params (shift,rotate)
    x0=[0,0,0;0,0,0]; % start with null transform: shift,rotate
    x2=x0';
    x2=x2(:)';
    lb=[-30 -30 -30 -pi/6 -pi/6 -pi/6];
    ub=[30 30 30 pi/6 pi/6 pi/6];
end

if flags.params == 9  % 9 params 
    x0=[0,0,0;0,0,0;1,1,1]; % start with null transform: shift;rot;scale
    x2=x0';
    x2=x2(:)';
    lb=[-30 -30 -30 -pi/6 -pi/6 -pi/6 .8 .8 .8 ];
    ub=[30 30 30 pi/6 pi/6 pi/6 1.2 1.2 1.2];
end
 
i_ter=1;    % initialize iteration counter
tic         % initialize timer

if flags.fit == 1 %use fmincon (per Elena)
    disp('Using fmincon')
    options=optimset('fmincon');% setting the optimization options
    options.LargeScale='off';
%     options.Display='iter';
%     options.TolPCG=0.01;
%     options.Display='iter';
    if flags.params == 3
        [x1,fval, flag, output] = fmincon(@xyz_MRI_fit3,x2,A,b,Aeq,beq,lb,ub,[],options);
    end
    if flags.params == 6
        
        [x1,fval, flag, output] = fmincon(@xyz_MRI_fit6,x2,A,b,Aeq,beq,lb,ub,[],options);
    end
    if flags.params==9
        options.TolX=.1;    %[.1, .1, .1, .1, .1, .1, .1, .1, .1];
        [x1,fval, flag, output] = fmincon(@xyz_MRI_fit9,x2,A,b,Aeq,beq,lb,ub,[],options);
    end
end

if flags.fit == 2
    options=optimset('lsqnonlin');% setting the optimization options
    options.LevenbergMarquardt='on';
%     options.LargeScale='off';
%     options.Display='iter';
%     options.TolPCG=0.01;
    disp('Using lsqnonlin (L-M)')
    if flags.params == 3
        [x1,fval, residual, flag, output] = lsqnonlin(@xyz_MRI_fit3,x2,lb,ub,options);
    end
    if flags.params == 6
        [x1,fval, residual, flag, output] = lsqnonlin(@xyz_MRI_fit6,x2,lb,ub,options);
    end
    if flags.params == 9
        [x1,fval, residual, flag, output] = lsqnonlin(@xyz_MRI_fit9,x2,lb,ub,options);
    end
end

if flags.fit == 3 % use "fminsearch"
    options = optimset('fminsearch');    % 'Display','Final','TolX',.01);
%     options.Display='iter';
    disp('Using fminsearch')
    if flags.params == 3
        [x1,fval,exitflag,output]=fminsearch(@xyz_MRI_fit3,x2,options);
    end
    if flags.params == 6
        [x1,fval,exitflag,output]=fminsearch(@xyz_MRI_fit6,x2,options);
    end
    if flags.params == 9
        [x1,fval,exitflag,output]=fminsearch(@xyz_MRI_fit9,x2,options);
    end
end

fit_time=toc;

% define output file names
%
xform_str=num2str(flags.params);
opt_str=num2str(flags.fit);
TRxyz_file=[in_path s_fold 'S' subj file_type '_em' xform_str '_O' opt_str '.mat'];
hist_fig_file=[in_path s_fold 'S' subj file_type '_em' xform_str '_O' opt_str '_history'];
elpreg_file=[in_path s_fold 'S' subj file_type '_em' xform_str '_O' opt_str '_elpreg'];

% plot fit history
%
annot_str=['S' subj file_type '_em' xform_str '_O' opt_str];
fig1=figure('Position',[0.05*scrsz(3) 0.025*scrsz(4) 0.9* scrsz(3) 0.9*scrsz(4)],...
    'Visible','on');
annotation('textbox',[.25 .95 .5 .05],'EdgeColor','none',...
           'String',annot_str,'Interpreter','none',...
           'FontSize',12,'HorizontalAlignment','center');
for i_plot=1:4
    subplot(2,2,1)
    plot(track_dist)
    xlim([1 length(track_dist)])
    ylim([0 5])
    title('RESIDUAL ERROR')
    subplot(2,2,2)
    plot(trackx0(:,1:3))
    xlim([1 length(track_dist)])
    if flags.params ==3;title('SCALE');else title('SHIFT');end
    if flags.params >3
        subplot(2,2,3)
        plot(trackx0(:,4:6))
        xlim([1 length(track_dist)])
        title('ROTATION')
        if flags.params >6
            subplot(2,2,4)
            plot(trackx0(:,7:9))
            xlim([1 length(track_dist)])
            title('SCALE')
        end
    end
end
% drawnow
print('-dtiff',hist_fig_file)

% Write fit results to screen
%
TRtxt_file=[in_path s_fold 'S' subj file_type '_em' xform_str '_O' opt_str '.txt'];
fid=1;
% fprintf(fid,'Writing to %s\n',TRtxt_file);
fprintf(fid,'Residual error = %5.2f mm.\n',fval);
fprintf(fid,'Best fit xform = \n');
fprintf(fid,'%8.4f  %8.4f  %8.4f\n',x1);
fprintf(fid,'Fit time = %5.2f\n\n',fit_time);

% Write fit results to text file
%
TRtxt_file=[in_path s_fold 'S' subj file_type '_em' xform_str '_O' opt_str '.txt'];
fid=fopen(TRtxt_file,'w');
fprintf(fid,'%s\n',TRtxt_file);
fprintf(fid,'Residual error = %5.2f mm.\n',fval);
fprintf(fid,'Best fit xform = \n');
fprintf(fid,'%8.4f  %8.4f  %8.4f\n',x1);
fprintf(fid,'Fit time = %5.2f\n',fit_time);
%output.
%        iterations: 65
%         funcCount: 902
%          stepsize: 0.0081
%      cgiterations: []
%     firstorderopt: []
%         algorithm: 'medium-scale: Levenberg-Marquardt, line-search'
%           message: [1x78 char]
% out_cell=struct2cell(output);
% for i=1:2;fprintf(fid,'%s\n',char(out_cell(i,:)));end
fclose(fid);

% 5) display initial fit
%
if flags.plot == 2
%     disp('Plotting initial registration')
    fig2=figure('Position',[0.05*scrsz(3) 0.025*scrsz(4) 0.9* scrsz(3) 0.9*scrsz(4)],...
        'Visible','on');
    subplot(2,3,1)
    plot_MRI_xyz(vert1,points_xyz,ClosestPts,'Fiducial (.tai) Alignment',MR_th)
    view(-90,0);
    subplot(2,3,2)
    plot_MRI_xyz(vert1,points_xyz,ClosestPts,xyz_file,MR_th)
    view(-90,90);
    subplot(2,3,3)
    plot_MRI_xyz(vert1,points_xyz,ClosestPts,'',MR_th)
    view(0,0);
%     drawnow

% 7) display transformed fit
%
%     disp('Plotting coregistered points')
%     figure(fig1)
    subplot(2,3,4)
    plot_MRI_xyz(vert1,TRxyz,ClosestPts,'Best Fit',MR_th)
    view(-90,0);
    subplot(2,3,5)
    plot_MRI_xyz(vert1,TRxyz,ClosestPts,TRxyz_file,MR_th)
    view(-90,90);
    subplot(2,3,6)
    plot_MRI_xyz(vert1,TRxyz,ClosestPts,'',MR_th)
    view(0,0);
%     drawnow
    print('-dtiff',elpreg_file)
%     disp(['Writing ' elpreg_file])
%     close(fig2)
end

% save the original & transformed elp coords
%
clear xyz
disp(['Saving ' TRxyz_file])
xyz=TRxyz;
save(TRxyz_file,'xyz','fit_time','x1','fval','output','trackx0','-mat');

% 8) write transformed elp+acpc.BRIK & HEAD
%
